import os
import sys

import pandas

from func_401_information_entropy import Ent

sys.path.append(os.pardir)


def Gain(D: pandas.DataFrame, a: str, final_tag: str):
    """
    英文名：information gain
    中文名：信息增益
    作用：计算出属性a对样本集D进行划分所获得的“信息增益”，可用于决策树的划分属性选择
    特性：
        信息增益越大，则意味着使用属性a来进行划分所获得的“纯度提升越大”

    :param D: list, 二维表格
    :param final_tag: 最终标签，如'好瓜'
    :param a: str, 某一属性，如'色泽'
    """
    Sigma = 0
    av = D[a].value_counts()  # 属性a的所有可能集合及个数
    for v in av.index:
        Sigma += (av[v] / len(D)) * Ent(D=D.loc[D[a] == v][final_tag])

    return Ent(D=D[final_tag]) - Sigma


if __name__ == '__main__':
    from table_401_watermelon_set_20 import df
    r = Gain(D=df, a='色泽', final_tag='好瓜')
    print(r)
